#include "udt_interface.h"
#include <memory.h>
#include <stdlib.h>
#include <stdio.h>


void udt_testSingleItem(UDTDataItem *pItem, UDTTreeLink pRoot, double *labelProb, int labelCount)
{
	int i;
	double labelProbBuf[MAX_LABEL];
    double prob;

	if (pRoot->testAttrib == -1)  {
		memcpy(labelProb, pRoot->labelProb, sizeof(double)*labelCount);

/*		printf("Test Leaf\n");
		for (i = 0; i < labelCount; ++i)
			printf("%lf ", labelProb[i]);
		printf("\n");
*/		
		return;
	};

	memset(labelProb, 0, sizeof(double)*labelCount);
	
	if (pItem->attrib[pRoot->testAttrib][0] < pRoot->splitPoint && pRoot->pLeft) {
		udt_testSingleItem(pItem, pRoot->pLeft, labelProbBuf, labelCount);

		if (pItem->attrib[pRoot->testAttrib][1] <= pRoot->splitPoint)
			prob = 1;
		else
			prob = (pRoot->splitPoint - pItem->attrib[pRoot->testAttrib][0])
				 / (pItem->attrib[pRoot->testAttrib][1] - pItem->attrib[pRoot->testAttrib][0]);

		for (i = 0; i < labelCount; ++i) { 
			labelProb[i] = labelProbBuf[i] * prob;
//			printf("%lf ", labelProb[i]);
		};
//		printf("\n");
	};

	if (pItem->attrib[pRoot->testAttrib][1] > pRoot->splitPoint && pRoot->pRight) {
		udt_testSingleItem(pItem, pRoot->pRight, labelProbBuf, labelCount);

		if (pItem->attrib[pRoot->testAttrib][0] >= pRoot->splitPoint)
			prob = 1;
		else
			prob = (pItem->attrib[pRoot->testAttrib][1] - pRoot->splitPoint)
				 / (pItem->attrib[pRoot->testAttrib][1] - pItem->attrib[pRoot->testAttrib][0]);

		for (i = 0; i < labelCount; ++i) 
			labelProb[i] += labelProbBuf[i] * prob;
	};
}; 

void _random_shuffle(int *data, int count)
{
    int n, k, temp;
    for (n = count; n > 1; n--) {
        k = rand() % n;
        temp = data[k];
        data[k] = data[n-1];
        data[n-1] = temp;
    };
}; 

void _genKFold(int dataCount, int *dataFold, int k)
{
    int i, j, last;
    int dataMap[MAX_ITEM];
    for (i = 0; i < dataCount; ++i)
        dataMap[i] = i;
    _random_shuffle(dataMap, i);
    last = 0;
    for (i = 0; i < k-1; ++i) {
       for (j = last; j < last + dataCount / k; ++j)
           dataFold[dataMap[j]] = i;
       last += dataCount / k;
    };
    for (j = last; j < dataCount; ++j)
        dataFold[dataMap[j]] = k-1;     
};    

double udt_testDataSet(UDTDataLinkSet *pLink, UDTTreeLink pRoot)
{
    int i, j, index, correct = 0;
	double labelProb[MAX_LABEL];    
    double max;
    for (i = 0; i < pLink->linkCount; ++i) {
        udt_testSingleItem(pLink->items[i], pRoot, labelProb, pLink->labelCount);
        max = 0;
        for (j = 0; j < pLink->labelCount; ++j)
            if (labelProb[j] > max) {
                index = j;
                max = labelProb[j];
            };
//		printf("predict: %d, label %d\n", index, pLink->items[i]->label);
        if (index == pLink->items[i]->label)
            correct++;
    };
    return (correct+0.0)/pLink->linkCount;
};    

void udt_testDataSet_File(UDTDataLinkSet *pLink, UDTTreeLink pRoot, char *fileName, int labelCount)
{
    int i, j, index, correct = 0;
	double labelProb[MAX_LABEL];    
    double max;
    FILE *fout = fopen(fileName, "w");

    for (i = 0; i < pLink->linkCount; ++i) {

        udt_testSingleItem(pLink->items[i], pRoot, labelProb, labelCount);
        max = 0;
        for (j = 0; j < labelCount; ++j)
            if (labelProb[j] > max) {
                index = j;
                max = labelProb[j];
            };
		fprintf(fout, "%d\n", index+1);
    };

    fclose(fout);
};    



double udt_testKFold(UDTDataLinkSet *pLink, int k)
{
    int i, j;
    int* dataFold = (int*)malloc(sizeof(int)*MAX_ITEM);
    int* attribUsed = (int*)malloc(sizeof(int)*MAX_ATTRIB);
    UDTDataLinkSet trainSet, testSet;
    UDTTreeLink pRoot;
    double precision, sum_precision;


	testSet.attribCount = trainSet.attribCount = pLink->attribCount;	
    testSet.labelCount = trainSet.labelCount = pLink->labelCount;
    testSet.linkCount = trainSet.linkCount = 0;


    _genKFold(pLink->linkCount, dataFold, k);


    sum_precision = 0;

    for (i = 0; i < k; ++i) {
		
		printf("Testing... Fold %d as test set\n", i);

    	testSet.linkCount = trainSet.linkCount = 0;			
        for (j = 0; j < pLink->linkCount; ++j)
            if (dataFold[j] == i) 
                testSet.items[testSet.linkCount++] = pLink->items[j];
            else
                trainSet.items[trainSet.linkCount++] = pLink->items[j];                       


		printf("Test Set: %d, Train Set: %d\n", testSet.linkCount, trainSet.linkCount);

        pRoot = (UDTTreeLink)malloc(sizeof(UDTTreeObject));

        memset(attribUsed, 0, sizeof(int)*MAX_ATTRIB);

        max_trained_depth = 0;
		udt_train(&trainSet, pRoot, attribUsed, 0);
        printf("Trained, max depth %d\n", max_trained_depth);

        precision = udt_testDataSet(&testSet, pRoot);
		sum_precision += precision;
    
		printf("Precision: %lf\n", precision); 

        udt_freeTree(pRoot);
    };

    free(dataFold);
    free(attribUsed);
    return sum_precision / k;
};



